package org.dxy.ai.config;


import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dxy.ai.repository.RedisChatMemoryStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.StringRedisTemplate;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.time.Duration;
import java.util.List;

/**
 * @ description:
 * @ author: dxy
 * @ createTime: 2025/10/4 23:41
 */
@Slf4j
@Configuration
@RequiredArgsConstructor
public class RagConfig {

    private final RedisChatMemoryStore redisChatMemoryStore;
    private final EmbeddingModel embeddingModel;
    private final RedisEmbeddingStore redisEmbeddingStore;
    private final StringRedisTemplate stringRedisTemplate;
    private static final String EMBEDDING_INIT_FLAG = "embeddingStore:initialized";


    @Bean
    public ChatMemory chatMemory() {
        return MessageWindowChatMemory.builder()
                .maxMessages(20)
                .build();
    }

    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return new ChatMemoryProvider() {
            @Override
            public ChatMemory get(Object memoryId) {
                return MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(20)
                        .chatMemoryStore(redisChatMemoryStore)
                        .build();
            }
        };
    }


    private static boolean initialized = false;

    @Bean("myStore")
    public EmbeddingStore<TextSegment> embeddingStore() {
        // 检查 Redis 标志
        Boolean acquired = stringRedisTemplate.opsForValue().setIfAbsent(EMBEDDING_INIT_FLAG, "true", Duration.ofHours(1));

        if (Boolean.TRUE.equals(acquired)) {
            // 执行一次向量化 + 入库
            List<Document> documents = ClassPathDocumentLoader.loadDocuments("content");
            EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                    .embeddingStore(redisEmbeddingStore)
                    .embeddingModel(embeddingModel)
                    .build();
            ingestor.ingest(documents);
            // 再确认标志持久化
            stringRedisTemplate.opsForValue().set(EMBEDDING_INIT_FLAG, "true");
        }

        return redisEmbeddingStore;
    }

    @Bean
    public ContentRetriever contentRetriever(){
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(redisEmbeddingStore)
                .embeddingModel(embeddingModel)
                .minScore(0.5) // 最小的余弦相似度
                .maxResults(3)
                .build();
    }

}
